package com.kdtree;

import com.jsurf.InterestPoint;

import java.util.Arrays;

public class KDTree {
	private static final String LINE = "no line";

	public static KDNode build(InterestPoint[] features, int n) {
		KDNode kdRoot;

		if(features == null || n <= 0) {
			System.err.format("Warning: (KDTree) build(): no features, %s, line %d\n", "KDTree", LINE);
			return null;
		}

		kdRoot = kdNodeInit(features, 0, n);
		expandKDNodeSubtree(kdRoot);

		return kdRoot;
	}

	public static int kdTreeBbfKnn(KDNode kdRoot, InterestPoint feature, int k,
	                               InterestPoint[] neighbors, int maxNNChecks) {
		KDNode expl;
		MinPriQueue minPQ;
		InterestPoint treeFeature, nbrs[];
		BBFData bbfData;
		int i, t = 0, n = 0;

		if(neighbors == null || feature == null || kdRoot == null) {
			System.err.format("Warning: null pointer error, %s, line %d\n", "KDTree", LINE);
			return -1;
		}

		nbrs = new InterestPoint[k];
		minPQ = new MinPriQueue();
		minPQ.insert(kdRoot, 0);
		boolean fail = false;
		while(minPQ.n > 0 && t < maxNNChecks) {
			expl = (KDNode)minPQ.extractMin();
			if(expl == null) {
				System.err.format("Warning: PQ unexpectedly empty, %s line %d\n", "KDTree", LINE);
				fail = true;
				break;
			}

			expl = exploreToLeaf(expl, feature, minPQ);
			if(expl == null) {
				System.err.format("Warning: PQ unexpectedly empty, %s line %d\n", "KDTree", LINE);
				fail = true;
				break;
			}

			for(i = 0; i < expl.n; i++) {
				treeFeature = expl.features[i];
				bbfData = new BBFData();
				if(bbfData == null) {
					System.err.format("Warning: unable to allocate memory, %s line %d\n", "KDTree", LINE);
					fail = true;
					break;
				}
				bbfData.oldData = treeFeature.featureData;
				bbfData.d = InterestPoint.descriptorSquaredDistance(feature, treeFeature);
				treeFeature.featureData = bbfData;
				n += insertIntoNeighborArray(treeFeature, nbrs, n, k);
			}
			t++;
		}

		if(!fail) {
			for(i = 0; i < n; i++ ) {
				bbfData = (BBFData)(nbrs[i].featureData);
				nbrs[i].featureData = bbfData.oldData;
			}
			for (int j = 0; j < nbrs.length; j++) {
				neighbors[j] = nbrs[j];
			}
			return n;
		} else {
			for(i = 0; i < n; i++) {
				bbfData = (BBFData)(nbrs[i].featureData);
				nbrs[i].featureData = bbfData.oldData;
			}
			return -1;
		}
	}

//	int kdTreeBbfSpatialKnn(KDNode kdRoot, InterestPoint feature,
//	                        int k, InterestPoint[] neighbors, int maxNNChecks,
//	                        Rect rect, int model)
//	{
//		InterestPoint all_nbrs[] = null, sp_nbrs[];
//		Point pt;
//		int i, n, t = 0;
//
//		n = kdTreeBbfKnn(kdRoot, feature, maxNNChecks, all_nbrs, maxNNChecks);
//		sp_nbrs = new InterestPoint[k]; //calloc( k, sizeof( InterestPoint[] ) );
//		for( i = 0; i < n; i++ )
//		{
//			if( model > 0 )
//				pt = all_nbrs[i].mdl_pt;
//			else
//				pt = all_nbrs[i].img_pt;
//
//			if( withinRect(pt, rect) > 0 )
//			{
//				sp_nbrs[t++] = all_nbrs[i];
//				if( t == k )
//					break;
//			}
//		}
//
//		neighbors = Arrays.copyOf(sp_nbrs, sp_nbrs.length);
//		return t;
//	}
//
//	void release(KDNode kdRoot)
//	{
//		if(kdRoot == null) {
//			return;
//		}
//		release(kdRoot.kdLeft);
//		release(kdRoot.kdRight);
//	}

	private static KDNode kdNodeInit(InterestPoint[] features, int from, int n) {
		KDNode kdNode;

		kdNode = new KDNode();
		kdNode.keyIndex = -1;
		kdNode.features = Arrays.copyOfRange(features, from, from + n);
		kdNode.n = n;

		return kdNode;
	}

	private static void expandKDNodeSubtree(KDNode kdNode) {
		/* base case: leaf node */
		if(kdNode.n == 1 || kdNode.n == 0) {
			kdNode.leaf = 1;
			return;
		}

		assignPartitionKey(kdNode);
		partitionFeatures(kdNode);

		if(kdNode.kdLeft != null) {
			expandKDNodeSubtree(kdNode.kdLeft);
		}
		if(kdNode.kdRight != null) {
			expandKDNodeSubtree(kdNode.kdRight);
		}
	}

	private static void assignPartitionKey(KDNode kdNode) {
		InterestPoint[] features;
		double keyValue, x, mean, var, varMax = 0;
		double tmp[];
		int d, n, i, j, keyIndex = 0;

		features = kdNode.features;
		n = kdNode.n;
		d = features[0].DESC_SIZE;

		// partition key index is that along which descriptors have most variance
		for(j = 0; j < d; j++) {
			mean = var = 0;
			for(i = 0; i < n; i++) {
				mean += features[i].getDescriptor()[j];
			}
			mean /= n;
			for(i = 0; i < n; i++) {
				x = features[i].getDescriptor()[j] - mean;
				var += x * x;
			}
			var /= n;

			if(var > varMax) {
				keyIndex = j;
				varMax = var;
			}
		}

		// partition key value is median of descriptor values at keyIndex
		tmp = new double[n]; // calloc( n, sizeof( double ) );
		for(i = 0; i < n; i++) {
			tmp[i] = features[i].getDescriptor()[keyIndex];
		}
		keyValue = medianSelect(tmp, n);
//		free( tmp );

		kdNode.keyIndex = keyIndex;
		kdNode.keyValue = keyValue;
	}

	static double medianSelect(double array[], int n) {
		return rankSelect(array, n, (n - 1) / 2);
	}

	static double rankSelect(double array[], int n, int r) {
		double tmp[], med;
		int gr5, grTot, remaining, i, j;

		// base case
		if(n == 1) {
			return array[0];
		}

		// divide array into groups of 5 and sort them
		gr5 = n / 5;
		grTot = (int)Math.ceil( n / 5.0 );
		remaining = n % 5;
//		tmp = Arrays.copyOf(array,5);
		for(i = 0; i < gr5; i++) {
//			insertion_sort( tmp, 5 );
			Arrays.sort(array, i * 5, Math.min(i * 5 + 5, array.length));
//			tmp += 5;
		}
//		insertion_sort( tmp, remaining );

		// recursively find the median of the medians of the groups of 5
		tmp = new double[grTot]; //calloc( grTot, sizeof( double ) );
		for(i = 0, j = 2; i < gr5; i++, j += 5) {
			tmp[i] = array[j];
		}
		if(remaining > 0) {
			tmp[i++] = array[n - 1 - remaining /2];
		}
		med = rankSelect(tmp, i, (i - 1) / 2);
//		free( tmp );

		// partition around median of medians and recursively select if necessary
		j = partitionArray(array, n, med);
		if(r == j) {
			return med;
		}
		else if(r < j) {
			return rankSelect(array, j, r);
		}
		else {
			return rankSelect(Arrays.copyOfRange(array, j + 1, n), (n - j - 1), (r - j - 1));
		}
	}

//	void insertionSort(double array[], int n)
//	{
//		double k;
//		int i, j;
//
//		for( i = 1; i < n; i++ )
//		{
//			k = array[i];
//			j = i-1;
//			while( j >= 0  &&  array[j] > k )
//			{
//				array[j+1] = array[j];
//				j -= 1;
//			}
//			array[j+1] = k;
//		}
//	}

	static int partitionArray(double array[], int n, double pivot)
	{
		double tmp;
		int p = 0, i, j;

		i = -1;
		for(j = 0; j < n; j++) {
			if( array[j] <= pivot ) {
				tmp = array[++i];
				array[i] = array[j];
				array[j] = tmp;
				if(array[i] == pivot) {
					p = i;
				}
			}
		}
		array[p] = array[i];
		array[i] = pivot;

		return i;
	}

	private static void partitionFeatures(KDNode kdNode) {
		InterestPoint features[], tmp;
		double keyValue;
		int n, keyIndex, p = 0, i, j = -1;

		features = kdNode.features;
		n = kdNode.n;
		keyIndex = kdNode.keyIndex;
		keyValue = kdNode.keyValue;
		for(i = 0; i < n; i++) {
			if( features[i].getDescriptor()[keyIndex] <= keyValue ) {
				tmp = features[++j];
				features[j] = features[i];
				features[i] = tmp;
				if(features[j].getDescriptor()[keyIndex] == keyValue) {
					p = j;
				}
			}
		}
		tmp = features[p];
		features[p] = features[j];
		features[j] = tmp;

		/* if all records fall on same side of partition, make node a leaf */
		if(j == n - 1) {
			kdNode.leaf = 1;
			return;
		}

		kdNode.kdLeft = kdNodeInit(features, 0, j + 1);
		kdNode.kdRight = kdNodeInit(features, j + 1, n - j - 1);
	}

	private static KDNode exploreToLeaf(KDNode kdNode, InterestPoint feature, MinPriQueue minPQ) {
		KDNode unexpl, expl = kdNode;
		double keyValue;
		int keyIndex;

		while(expl != null  &&  expl.leaf == 0) {
			keyIndex = expl.keyIndex;
			keyValue = expl.keyValue;

			if(keyIndex >= feature.DESC_SIZE) {
				System.err.format("Warning: comparing imcompatible descriptors, %s line %d\n", "KDTree", LINE);
				return null;
			}
			if(feature.getDescriptor()[keyIndex] <= keyValue) {
				unexpl = expl.kdRight;
				expl = expl.kdLeft;
			}
			else {
				unexpl = expl.kdLeft;
				expl = expl.kdRight;
			}

			if(minPQ.insert(unexpl, (int) Math.abs(keyValue - feature.getDescriptor()[keyIndex])) > 0) {
				System.err.format("Warning: unable to insert into PQ, %s, line %d\n", "KDTree", LINE);
				return null;
			}
		}

		return expl;
	}

	public static int insertIntoNeighborArray(InterestPoint feature, InterestPoint[] neighbors, int n, int k) {
		BBFData fdata, ndata;
		double dn, df;
		int i, ret = 0;

		if(n == 0) {
			neighbors[0] = feature;
			return 1;
		}

		/* check at end of array */
		fdata = (BBFData) feature.featureData;
		df = fdata.d;
		ndata = (BBFData) neighbors[n - 1].featureData;
		dn = ndata.d;
		if(df >= dn) {
			if(n == k) {
				feature.featureData = fdata.oldData;
				return 0;
			}
			neighbors[n] = feature;
			return 1;
		}

		/* find the right place in the array */
		if(n < k) {
			neighbors[n] = neighbors[n - 1];
			ret = 1;
		}
		else {
			neighbors[n - 1].featureData = ndata.oldData;
		}
		i = n - 2;
		while(i >= 0) {
			ndata = (BBFData)neighbors[i].featureData;
			dn = ndata.d;
			if(dn <= df) {
				break;
			}
			neighbors[i + 1] = neighbors[i];
			i--;
		}
		i++;
		neighbors[i] = feature;

		return ret;
	}

//	public static int withinRect(Point pt, Rect rect) {
//		if(pt.x < rect.x || pt.y < rect.y) {
//			return 0;
//		}
//		if(pt.x > rect.x + rect.width || pt.y > rect.y + rect.height) {
//			return 0;
//		}
//		return 1;
//	}
}
